clear; clc;

I = im2double(imread('cameraman.tif'));
params.a = 1; params.b = 2*pi; params.c = 11.5;
params.x0 = 0.4; params.y0 = 0.8; params.z0 = 0.6;
params.m = 100; params.n = 500; params.CR = 0.8;

[C, keys] = encrypt_image(I, params);

% 显示原图和加密图
figure;
subplot(1,2,1); imshow(I); title('Original Image');
subplot(1,2,2); imshow(C); title('Encrypted Image');

% 保存加密图像为PNG文件，自动转换为uint8格式
imwrite(C, 'Encrypted_Image.png');
disp('Encrypted Image saved as Encrypted_Image.png');

function [C, keys] = encrypt_image(I, params)
    % 输入：
    % I      - 原始灰度图像，double类型，范围[0,1]
    % params - 结构体参数，包含3D-SIMM参数，迭代次数，压缩率等
    %
    % 输出：
    % C      - 加密图像 uint8矩阵
    % keys   - 结构体，保存密钥参数，供解密使用

    % Step 1: DWT稀疏化，置零细节系数，恢复同尺寸稀疏矩阵
    [cA, cH, cV, cD] = dwt2(I, 'haar');
    cH(:) = 0; cV(:) = 0; cD(:) = 0;
    I1 = idwt2(cA, cH, cV, cD, 'haar');
    I1_vec = I1(:);

    % 补齐长度为2的幂
    len_orig = length(I1_vec);
    len = 2^nextpow2(len_orig);
    if len > len_orig
        I1_vec(len_orig+1:len) = 0;
    end

    % 计算采样点数M和边长side_len（完全平方数）
    M_approx = floor(len * params.CR);
    side_len = floor(sqrt(M_approx));
    M = side_len^2;
    if M == 0
        error('采样点数 M 计算为0，压缩率过低或图像尺寸太小');
    end

    % Step 2: 生成混沌序列X，长度m+M，取后M
    len_X = params.m + M;
    [X, ~, ~] = gen_3DSIMM_seq(params.a, params.b, params.c, ...
                               params.x0, params.y0, params.z0, len_X);
    X = X(params.m+1:end);

    % Step 3: 根据X排序得到采样行索引
    [~, idx] = sort(X, 'descend');
    selected_rows = idx(1:M);

    % Step 4: 快速Walsh-Hadamard采样
    I2_vec = fwht_sampling(I1_vec, selected_rows);

    % Step 5: 重塑为二维压缩图像
    I2 = reshape(I2_vec(1:M), side_len, side_len);

    % Step 6: 归一化量化到[0,255]
    I_min = min(I2(:));
    I_max = max(I2(:));
    if I_max == I_min
        error('归一化失败，最大值等于最小值');
    end
    I2_norm = (I2 - I_min) / (I_max - I_min);
    I2q = uint8(floor(I2_norm * 255));

    % Step 7: 生成置乱序列，长度m+M*M，取后部分
    len_scramble = params.m + M;
    [X_scr, Y_scr, ~] = gen_3DSIMM_seq(params.a, params.b, params.c, ...
                                       params.x0, params.y0, params.z0, len_scramble);
    X_scr = X_scr(params.m+1:end);
    Y_scr = Y_scr(params.m+1:end);

    a_scr = floor(abs(X_scr(1) * 1e6));
    b_scr = floor(abs(Y_scr(1) * 1e6));

    % Step 8: Arnold二维置乱（示范用10次迭代）
    I3 = arnold_2d_scramble(I2q, 10);

    % Step 9: 矩阵转置
    I3 = I3';

    % Step 10: 生成扩散密钥序列
    len_diff = params.n + M;
    [~, y_diff, z_diff] = gen_3DSIMM_seq(params.a, params.b, params.c, ...
                                         params.x0, params.y0, params.z0, len_diff);
    y_diff = y_diff(params.n+1:end);
    z_diff = z_diff(params.n+1:end);

    S1 = mod(floor(y_diff * 2^10), 256);
    S2 = mod(floor(z_diff * 2^10), 256);

    % Step 11: 双向扩散加密
    C = diffusion_double_fixed(I3, S1, S2);

    % 返回密钥信息
    keys.Phi_rows = selected_rows;
    keys.I_min = I_min;
    keys.I_max = I_max;
    keys.side_len = side_len;
    keys.a_scr = a_scr;
    keys.b_scr = b_scr;
    keys.S1 = S1;
    keys.S2 = S2;
end

%% 辅助函数

function y = fwht_sampling(x, rows)
    N = length(x);
    if bitand(N, N-1) ~= 0
        error('输入向量长度必须是2的幂');
    end
    Hx = fwht(x) / sqrt(N);
    y = Hx(rows);
end

function [X, Y, Z] = gen_3DSIMM_seq(a, b, c, x0, y0, z0, len)
    X = zeros(1,len); Y = zeros(1,len); Z = zeros(1,len);
    X(1)=x0; Y(1)=y0; Z(1)=z0;
    for i=1:len-1
        [X(i+1), Y(i+1), Z(i+1)] = SIMM(X(i), Y(i), Z(i), a, b, c);
    end
end

function [x_next, y_next, z_next] = SIMM(x, y, z, a, b, c)
    x_next = a*sin(b*z)*sin(c/x);
    y_next = a*sin(b*x_next)*sin(c/y);
    z_next = a*sin(b*y_next)*sin(c/z);
end

function I3 = arnold_2d_scramble(I2q, iter)
    [N, ~] = size(I2q);
    I3 = I2q;
    for k = 1:iter
        temp = zeros(N,N,'uint8');
        for x = 0:N-1
            for y = 0:N-1
                new_x = mod(x + y, N);
                new_y = mod(x + 2*y, N);
                temp(new_x+1, new_y+1) = I3(x+1, y+1);
            end
        end
        I3 = temp;
    end
end

function C = diffusion_double_fixed(I3, S1, S2)
    [M, N] = size(I3);
    len = M*N;
    I_vec = reshape(I3, [], 1);
    C_vec = zeros(len, 1, 'uint16');
    
    % 前向扩散
    C_vec(1) = mod(double(I_vec(1)) + double(S1(1)), 256);
    for i = 2:len
        C_vec(i) = mod(C_vec(i-1) + double(I_vec(i)) + double(S1(i)), 256);
    end
    
    % 逆向扩散
    for i = len-1:-1:1
        C_vec(i) = mod(C_vec(i) + C_vec(i+1) + double(S2(i)), 256);
    end
    
    C = uint8(C_vec);
    C = reshape(C, M, N);
end

